import json
import openai
from collections import deque
from datetime import datetime, timedelta


class LLM:
    def __init__(self):
        api_key = '<用户 API Key>'
        base_url = 'https://dashscope.aliyuncs.com/compatible-mode/v1'
        self.model = 'qwen-turbo-latest'
        self.client = openai.OpenAI(api_key=api_key, base_url=base_url)
    
    def generate_response(self, messages, tools=None, tool_choice=None):
        try:
            response = self.client.chat.completions.create(
                model=self.model,
                messages=messages,
                tools=tools,
                tool_choice=tool_choice,
                temperature=0.5,
                max_tokens=2000
            )
            return response.choices[0].message
        except Exception as e:
            return f'生成响应时出错：{str(e)}'

class MemoryManager:
    def __init__(self, max_history=5):
        self.history = deque(maxlen=max_history)  # 固定长度的对话历史队列
        self.data_cache = {}  # 数据集缓存字典
    
    def add_history(self, data):
        self.history.append(data)  # 添加对话记录

    def get_history(self):
        '''获取对话历史'''
        return list(self.history)

class WeatherAgent:
    '''天气查询智能体'''
    def __init__(self):
        self.llm = LLM()
        self.memory = MemoryManager()
        self.system_prompt = {
            'role': 'system',
            'content': '''你是一个天气查询助手，能够使用工具获取实时天气信息。请按以下规则处理：
1. 当用户询问天气时，先调用工具获取数据再回答
2. 若用户未明确日期，默认查询当天天气
3. 回答需包含地点、日期和天气信息
4. 如果没有地点，需要提醒用户输入地点
5.可能需要多次选择工具
6. 非天气问题请保持回答简洁'''
        }
        self.tools = [
            {
                'type': 'function',
                'function': {
                    'name': 'get_date',
                    'description': '获取具体查询日期',
                    'parameters': {
                        'type': 'object',
                        'properties': {
                            'date': {
                                'type': 'string',
                                'description': '输入日期',
                                'enum': ['今天', '明天', '后天']
                            }
                        },
                        'required': ['date']
                    }
                }
            },
            {
                'type': 'function',
                'function': {
                    'name': 'get_weather',
                    'description': '查询指定地点的实时天气',
                    'parameters': {
                        'type': 'object',
                        'properties': {
                            'location': {
                                'type': 'string',
                                'description': '城市名称，如：北京、上海'
                            },
                            'date': {
                                'type': 'string',
                                'description': '查询日期，格式为yyyy-mm-dd'
                            }
                        },
                        'required': ['location', 'date']
                    }
                }
            }
        ]
    
    
    def _handle_tool_call(self, tool_call):
        '''处理工具调用请求'''
        func_name = tool_call.function.name
        args = eval(tool_call.function.arguments)
        if func_name  == 'get_date':
            return self._get_date(args)
        elif func_name == 'get_weather':
            return self._get_weather(args)
        return None
    
    def _get_date(self, params):
        # 获取今天的日期
        today = datetime.now()
        today_str = today.strftime('%Y-%m-%d')
        # 获取明天的日期
        tomorrow = today + timedelta(days=1)
        tomorrow_str = tomorrow.strftime('%Y-%m-%d')
        # 获取后天的日期
        day_after_tomorrow = today + timedelta(days=2)
        after_tomorrow_str = day_after_tomorrow.strftime('%Y-%m-%d')
        date_str = params.get('date', '今天')
        if date_str == '明天':
            return {'date': tomorrow_str}
        elif date_str == '后天':
            return {'date': after_tomorrow_str}
        return {'date': today_str}

    def _get_weather(self, params):
        '''天气查询工具实现'''
        # 模拟数据，也可以替换成其他天气查询的 API 接口
        # 获取今天的日期
        today = datetime.now()
        today_str = today.strftime('%Y-%m-%d')
        # 获取明天的日期
        tomorrow = today + timedelta(days=1)
        tomorrow_str = tomorrow.strftime('%Y-%m-%d')
        # 获取后天的日期
        day_after_tomorrow = today + timedelta(days=2)
        after_tomorrow_str = day_after_tomorrow.strftime('%Y-%m-%d')
        
        weather_infos = {
            '北京': {today_str: 25, tomorrow_str: 26, after_tomorrow_str: 24},
            '上海': {today_str: 26, tomorrow_str: 27, after_tomorrow_str: 28},
            '广州': {today_str: 27, tomorrow_str: 28, after_tomorrow_str: 29},
        }
        location = params.get('location', 'default')
        if location not in  weather_infos.keys():
            return {'temp': '未查询到当前地区天气，当前仅支持北京、上海、广州的天气查询。'}
        date = params.get('date', 'default')
        if date not in [today_str, tomorrow_str, after_tomorrow_str]:
            return {'temp': '未查询到当前时间段天气，当前仅支持今天、明天、后天的天气查询。'}
        return {'temp': weather_infos[location][date]}

    def process_query(self, query):
        '''处理用户查询的完整流程'''
        # 构建消息历史
        messages = [self.system_prompt] + self.memory.get_history()
        messages.append({'role': 'user', 'content': query})
        self.memory.add_history({'role': 'user', 'content': query})

        # 第一轮调用：获取工具调用请求
        # 调用大语言模型（LLM）生成响应，传入消息列表和可用工具列表
        response = self.llm.generate_response(messages, tools=self.tools)
        # 检查响应中是否包含工具调用信息
        if response.tool_calls is None:
            # 如果不包含工具调用信息，将响应内容添加到内存的历史记录中，并返回结果
            self.memory.add_history({'role': 'assistant', 'content': response.content})
            return response.content
        # 初始化迭代次数
        n_iter = 0
        # 大模型可能多次调用工具
        # 当响应中包含工具调用信息时，进入循环处理
        while response.tool_calls is not None:
            n_iter = n_iter + 1
            # 保存工具调用请求
            # 将工具调用请求以字典形式添加到消息列表中
            messages.append({
                    'role': 'assistant',
                    'content': '',
                    'tool_calls': response.tool_calls
                })
            
            #  处理工具调用
            # 遍历响应中的每个工具调用请求
            for i in range(len(response.tool_calls)):
                print('调用工具与参数：', response.tool_calls[i].function.name, response.tool_calls[i].function.arguments)
                # 调用内部方法处理工具调用，并获取工具的响应结果
                tool_response = self._handle_tool_call(response.tool_calls[i])
                print('工具响应结果：', tool_response)
                # 将工具的响应结果以字典形式添加到消息列表中
                messages.append({
                'role': 'tool',
                'content': json.dumps(tool_response),
                'tool_call_id': response.tool_calls[i].id
                })
            # 根据再次调用工具结果再次调用大语言模型生成响应，传入更新后的消息列表和可用工具列表
            response = self.llm.generate_response(messages, tools=self.tools)
            # print(messages)
            # print(response)
            # 如果响应中不再包含工具调用信息，则保存历史并返回当前结果
            if response.tool_calls is None:  
                self.memory.add_history({'role': 'assistant', 'content': response.content})
                return response.content
            # 防止因反复调用工具陷入死循环
            # 最多连续调用3次，超出后不再调用工具，保存历史并返回当前结果
            if n_iter > 3:
                response = self.llm.generate_response(messages)
                self.memory.add_history({'role': 'assistant', 'content': response.content})
                return response.content
        

# 大模型回复测试
def llm_test():
    llm = LLM()
    messages = [{'role': 'user', 'content': '你好'}]
    print(llm.generate_response(messages))

# 模拟数据测试
def weather_data_test():
    # 获取今天的日期
    today = datetime.now()
    today_str = today.strftime('%Y-%m-%d')
    # 获取明天的日期
    tomorrow = today + timedelta(days=1)
    tomorrow_str = tomorrow.strftime('%Y-%m-%d')
    # 获取后天的日期
    day_after_tomorrow = today + timedelta(days=2)
    after_tomorrow_str = day_after_tomorrow.strftime('%Y-%m-%d')
    
    weather_infos = {
        '北京': {today_str: 25, tomorrow_str: 26, after_tomorrow_str: 24},
        '上海': {today_str: 26, tomorrow_str: 27, after_tomorrow_str: 28},
        '广州': {today_str: 27, tomorrow_str: 28, after_tomorrow_str: 29},
    }
    print(f'天气模拟数据\n{weather_infos}')



# 测试多轮对话
def main():
    '''主运行程序'''
    agent = WeatherAgent()
    queries = [
        '上海的天气怎么样？',
        '那明天呢？',
        '后天北京的气温多少？'
    ]
    for query in queries:
        print(f'[用户] {query}')
        response = agent.process_query(query)
        print(f'[助手] {response}')
        print('-' * 60)

if __name__ == '__main__':
    # 测试大模型
    # llm_test()

    # 模拟数据测试
    weather_data_test()

    # 多轮对话
    print('进入天气查询智能体多轮对话：\n')
    main()